import os
import cv2
import numpy as np
import pandas as pd

import cv2
import numpy as np
import cv2
import numpy as np

def get_brightness(grid):
    return np.mean(grid)

def calculate_angle(point1, point2):
    """计算两个点之间的角度"""
    return np.arctan2(point2[1] - point1[1], point2[0] - point1[0])

# 这个是我自己的奇思妙想。因为当时，虽然最亮的线大致是正确的，但它会有横向的波动（其实现在也有一些，只不过后续进行了一定的平滑）
# 横向波动的效果就是：最热航线突然从一边移到另一边，非常不连续且不符合实际。
# 所以我当时想了一个新的观点就是：直接设计一个导航点，让它不断在行动过程中，收集周围邻域的亮度信息，不断找到最亮的方向行驶。也就是说，我直接模拟了一个找寻信息素的蚂蚁，或者想要找寻最热航线的传播，有点像智能算法，让最热航线的绘制变得自然，符合常理。
# 但具体的思想，可能会有些复杂，在变换角度的时候遇到的困难也挺多。最后的效果只能说还不够好，所以我暂时放弃了这个想法的进一步深究，假以时日，说不定可以实现这个构想。
def analyze_hottest_route(image_path, grid_size=(100, 100), neighbor_size=3, alpha=0.4, beta=0.6):
    """
    读取保存的图像，进行网格化分割和亮度分析，并进行路径规划。
    image_path: 图像文件的路径。
    grid_size: 网格大小。
    neighbor_size: 邻域大小。
    alpha: 航标方向权重。
    beta: 亮度累计值权重。
    """
    img = cv2.imread(image_path)
    height, width, _ = img.shape
    grid_h, grid_w = grid_size
    h_step, w_step = height // grid_h, width // grid_w

    # 初步最热轨迹绘制
    initial_route = []
    initial_directions = []
    for i in range(0, height, h_step):
        max_brightness = 0
        brightest_point = (0, i)
        for j in range(0, width, w_step):
            grid = img[i:i + h_step, j:min(j + w_step, width)]
            brightness = get_brightness(grid)
            if brightness > max_brightness:
                max_brightness = brightness
                brightest_point = (j + w_step // 2, i + h_step // 2)
        if initial_route:
            initial_directions.append(calculate_angle(initial_route[-1], brightest_point))
        initial_route.append(brightest_point)

    # 边缘亮度收集
    edge_points = []
    for x in range(0, width, w_step):
        edge_points.append((x, 0))
        edge_points.append((x, height - h_step))
    for y in range(0, height, h_step):
        edge_points.append((0, y))
        edge_points.append((width - w_step, y))

    max_brightness = -np.inf
    start_point = None
    for point in edge_points:
        brightness = get_neighbor_brightness(img, point[0], point[1], w_step, h_step, neighbor_size)
        if brightness > max_brightness:
            max_brightness = brightness
            start_point = point

    current_point = start_point
    hottest_route = [current_point]

    # 计算全局亮度阈值
    global_brightness_threshold = np.mean(img) * (2 * neighbor_size + 1) ** 2

    # 导航点设置
    for i in range(1, len(initial_route)):
        max_value = -np.inf
        next_point = None
        for dx in range(-neighbor_size, neighbor_size + 1):
            for dy in range(-neighbor_size, neighbor_size + 1):
                nx, ny = current_point[0] + dx * w_step, current_point[1] + dy * h_step
                if 0 <= nx < width and 0 <= ny < height:
                    brightness = get_neighbor_brightness(img, nx, ny, w_step, h_step, neighbor_size)
                    direction = calculate_angle(current_point, (nx, ny))
                    initial_direction = initial_directions[i-1]
                    angle_diff = np.abs(direction - initial_direction)

                    # 计算总值
                    total_value = alpha * (1 - angle_diff / np.pi) + beta * brightness
                    if total_value > max_value:
                        max_value = total_value
                        next_point = (nx, ny)

        if next_point:
            hottest_route.append(next_point)
            current_point = next_point

            # # 判断是否达到终点条件
            # if max_value < global_brightness_threshold/4:
            #     break

    # 绘制最终最热航线
    for i in range(len(hottest_route) - 1):
        cv2.line(img, hottest_route[i], hottest_route[i + 1], (0, 255, 0), thickness=8)

    # 保存处理后的图像
    cv2.imwrite(image_path.replace('.png', '_hottest_route.png'), img)

def get_neighbor_brightness(img, x, y, w_step, h_step, neighbor_size):
    """
    计算指定点邻域的亮度。
    """
    brightness = 0
    count = 0
    for dx in range(-neighbor_size, neighbor_size + 1):
        for dy in range(-neighbor_size, neighbor_size + 1):
            nx, ny = x + dx * w_step, y + dy * h_step
            if 0 <= nx < img.shape[1] and 0 <= ny < img.shape[0]:
                grid = img[ny:ny + h_step, nx:nx + w_step]
                brightness += np.mean(grid)
                count += 1
    return brightness / count if count > 0 else 0

if __name__ == "__main__":
    grouped_dir = "./Data/Groups/CSJ/New_Human_Trajectories/MMSI_Cluster_Groups/MMSI_Patterns"
    detailed_dir = "./Data/Groups/CSJ/New_Human_Trajectories/MMSI_Cluster_Groups/Detailed_Trajectories"
    output_dir = "Data/Groups/CSJ/New_Human_Trajectories/MMSI_Cluster_Groups/Hottest_Plot_for_CV2"

    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    # 遍历所有MMSI模式组文件
    for group_file in os.listdir(grouped_dir):
        group_df = pd.read_csv(os.path.join(grouped_dir, group_file))
        save_path = os.path.join(output_dir, f"{group_file[:-4]}_heatmap.png")
        analyze_hottest_route(save_path)
